import sys, os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from PyQt5.QtCore import Qt, pyqtSignal, QObject
from PyQt5.QtGui import QColor, QBrush
from PyQt5.QtWidgets import (QApplication, QMainWindow, QAction, QFileDialog,
                             QTableWidget, QTableWidgetItem, QVBoxLayout,
                             QWidget, QToolBar, QStatusBar, QMessageBox,
                             QAbstractItemView, QSplitter,
                             QPushButton, QCheckBox, QGroupBox, QLabel)


class TableViewer(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("表格阅读器（左右分割版）")
        self.resize(2000, 1200)

        # ---------- 左右两列 ----------
        central = QWidget()  # 中央控件
        self.setCentralWidget(central)

        layout = QVBoxLayout(central)  # 创建一垂直布局器
        layout.setContentsMargins(0, 0, 0, 0)  # 设置这个布局器的大小

        splitter = QSplitter(Qt.Horizontal, central)  # 创建一个左右分割的布局，为顶层的Qsplitter
        layout.addWidget(splitter)  # 布局器中添加这个布局

        """左侧：只有按钮"""
        left_widget = QWidget()  # 创建一个QWidget()实例作为容器
        left_layout = QVBoxLayout(left_widget)  # 创建一个垂直的布局管理器，并设置left_widget作为该容器的父容器
        left_layout.setAlignment(Qt.AlignTop)  # 设置布局的对其方式

        # 添加一些控件
        grp = QGroupBox("表格展示")
        vbox = QVBoxLayout(grp)  # 创建一个垂直布局
        self.open_btn = QPushButton("打开文件", left_widget)  # 创建一个按钮控件PushButton
        self.open_btn.clicked.connect(self.open_file)  # 按钮的点击事件连接到槽函数
        vbox.addWidget(self.open_btn)  # 按钮添加到布局管理器中

        grp_2 = QGroupBox("表格内容定制")
        vbox_2 = QVBoxLayout(grp_2)
        self.color_check_box = QCheckBox("表格上色", left_widget)
        self.color_check_box.setChecked(False)
        self.color_check_box.stateChanged.connect(self.on_check_changed)
        vbox_2.addWidget(self.color_check_box)

        self.confirm_btn = QPushButton("确认", left_widget)
        vbox.addWidget(self.confirm_btn)

        vbox.addWidget(grp_2)

        left_layout.addWidget(grp)

        """右侧：原表格"""
        self.table = QTableWidget()

        splitter.addWidget(left_widget)  # 添加左侧的元素
        splitter.addWidget(self.table)  # 添加右侧的元素
        splitter.setStretchFactor(0, 1)  # 设置左侧占据空间的大小，左 1 份
        splitter.setStretchFactor(1, 2)  # 设置右侧占据空间的大小，右 2 份
        splitter.setSizes([200, 1200])

        # ---------- 菜单/工具栏/状态栏 ----------
        menubar = self.menuBar()  # 菜单栏
        file_menu = menubar.addMenu("文件(&F)")

        open_action = QAction("打开(&O)...", self)
        open_action.setShortcut("Ctrl+O")
        open_action.triggered.connect(self.open_file)
        file_menu.addAction(open_action)

        save_action = QAction("保存(&S)...", self)
        save_action.setShortcut("Ctrl+S")
        save_action.triggered.connect(self.save_file)
        file_menu.addAction(save_action)

        toolbar = QToolBar("主工具栏")  # 工具栏
        self.addToolBar(toolbar)
        toolbar.addAction(open_action)
        toolbar.addAction(save_action)

        self.status = QStatusBar()
        self.setStatusBar(self.status)
        self.status.showMessage("就绪")

    def open_file(self):
        path, _ = QFileDialog.getOpenFileName(
            self, "打开表格", "", "Excel (*.xlsx);;CSV (*.csv);;All Files (*)")
        if not path:
            return
        try:
            if path.lower().endswith(".csv"):
                df = pd.read_csv(path)
            else:
                df = pd.read_excel(path)
                df.set_index(df.columns[0], inplace=True)
        except Exception as e:
            QMessageBox.critical(self, "读取失败", str(e))
            return
        self.populate_table(df)
        self.status.showMessage(f"已加载：{os.path.basename(path)}")

    def save_file(self):
        if self.table.rowCount() == 0:
            QMessageBox.information(self, "提示", "表格为空，无需保存。")
            return
        path, _ = QFileDialog.getSaveFileName(
            self, "保存表格", "", "Excel (*.xlsx);;CSV (*.csv);;All Files (*)")
        if not path:
            return
        rows, cols = self.table.rowCount(), self.table.columnCount()
        header = [self.table.horizontalHeaderItem(c).text() for c in range(cols)]
        data = [[self.table.item(r, c).text() if self.table.item(r, c) else ""
                 for c in range(cols)] for r in range(rows)]
        df = pd.DataFrame(data, columns=header)
        try:
            if path.lower().endswith(".csv"):
                df.to_csv(path, index=False, encoding='utf-8-sig')
            else:
                df.to_excel(path, index=False)
            self.status.showMessage(f"已保存：{os.path.basename(path)}")
        except Exception as e:
            QMessageBox.critical(self, "保存失败", str(e))

    def populate_table(self, df):
        self.table.clear()
        self.table.setRowCount(len(df))
        self.table.setColumnCount(len(df.columns))
        self.table.setHorizontalHeaderLabels(df.columns.astype(str))
        self.table.setVerticalHeaderLabels(df.index.astype(str))
        for r in range(len(df)):
            for c in range(len(df.columns)):
                self.table.setItem(r, c, QTableWidgetItem(str(df.iat[r, c])))
        self.table.resizeColumnsToContents()
        self.table.setAlternatingRowColors(True)
        self.table.setEditTriggers(QAbstractItemView.DoubleClicked)

    def on_check_changed(self, state):
        if state == Qt.Checked:
            self.set_table_color()
        else:
            self.clear_table_color()

    def set_table_color(self):
        df = self.get_table_data()
        if df is None:
            return
        mn, mx = df.min().min(), df.max().max()
        if mx == mn:
            return
        cmap = plt.get_cmap('turbo')
        for r in range(self.table.rowCount()):
            for c in range(self.table.columnCount()):
                try:
                    val = float(self.table.item(r, c).text())
                    norm = (val - mn) / (mx - mn)
                    col = cmap(norm)
                    brush = QBrush(QColor(*(int(x * 255) for x in col[:3]), 128))
                    self.table.item(r, c).setBackground(brush)
                except ValueError:
                    pass
        self.table.update()  # 刷新表格显示

    def clear_table_color(self):
        for r in range(self.table.rowCount()):
            for c in range(self.table.columnCount()):
                self.table.item(r, c).setBackground(QBrush())
        self.table.update()  # 刷新表格显示

    def get_table_data(self):
        rows, cols = self.table.rowCount(), self.table.columnCount()
        header = [self.table.horizontalHeaderItem(c).text() for c in range(cols)]
        data = [[float(self.table.item(r, c).text()) if self.table.item(r, c) else ""
                 for c in range(cols)] for r in range(rows)]
        df = pd.DataFrame(data, columns=header)
        return df


if __name__ == "__main__":
    app = QApplication(sys.argv)
    win = TableViewer()
    win.show()
    sys.exit(app.exec_())
